@andygrove cc as this ticket directly related on Comet working on cloud instances

I had a hard time making DataFusion Comet work on cloud instances with 4GB memory per CPU core, partially because DataFusion is very likely to allocates more memory than reserved and run into OOM, or run into various kinds of memory reservation failures.

Yes, this feature is quite bug-prone, perhaps we should mark it as experimental to prevent someone to use it in production. Thank you so much for the efforts.

Here are my thoughts on the changes

1. Don't try_collect the result of merging all at once. We consume the merged batches one after another and reserve memory for each batch. Once the reservation fails we switch to "spill mode" and write all future batches into the spill file. This resolves the 2X memory allocation problem ("the second problem") reported by Memory account not adding up in SortExec #10073, as well as this comment: External sorting not working for (maybe only for string columns??) #12136 (comment)

3. Reserves more memory for ingested batches to leave some room for merging. This PR reserves 2X memory for each batch, this works for most of the queries in sort-tpch benchmark (all except Q5 and Q8). User still have to configure sort_spill_reservation_bytes when memory reserved is not big enough for merging. I don't think it is a good change but this is the only solution I can think of to compensate for the extra memory usage for the row representation of sorted columns.

The 2X memory problem is: specifying a query to run under 100M memory, the measured physical memory is 200M. (though the reality is even worse than 2X 🤦🏼 , see #14142)
This is caused by, when the first time OOM happens:

• there are already 1X batches in memory
• Then it will be sorted and merged at once, meaning original batches should hold until all sorted runs are generated (now memory footprint is already 2X)
• In the merging phase, additional col->row conversion consumes some extra memory (2X+)

I can get the high-level idea of the more conservative memory accounting in point 3, it is used to also account for merged batches, but I get lost in the memory estimation details in the implementation (especially is this 2X estimation amplification only for merged batches, or also intermediate Rows), could you explain with a concrete example how everything is calculated? (for example, there is a ExternalSorter with 100M memory budget, and it will consume 10 batches, each with 20M size, how memory estimation is calculated in each step)

2. shrink_to_fit every sorted batches reduce the memory footprint of sorted batches, otherwise sorted string arrays may take 2X the original space in the worst case, due to exponential growth of MutableBuffer for storing variable length binary values. shrink_to_fit is a no-op for primitive-type columns returned by take_arrays since they already have the right capacity, and benchmarking showed no significant performance regression for non-primitive types such as string arrays, so I think it is a good change. This resolves "the first problem" reported by Memory account not adding up in SortExec #10073.

I think the buffer resizing mechanism is not doubling each time, the default policy will allocate new constant size buffers https://docs.rs/arrow-array/54.1.0/src/arrow_array/builder/generic_bytes_view_builder.rs.html#120-122, so this change might not help

2. shrink_to_fit every sorted batches reduce the memory footprint of sorted batches, otherwise sorted string arrays may take 2X the original space in the worst case, due to exponential growth of MutableBuffer for storing variable length binary values. shrink_to_fit is a no-op for primitive-type columns returned by take_arrays since they already have the right capacity, and benchmarking showed no significant performance regression for non-primitive types such as string arrays, so I think it is a good change. This resolves "the first problem" reported by Memory account not adding up in SortExec #10073.

I think the buffer resizing mechanism is not doubling each time, the default policy will allocate new constant size buffers https://docs.rs/arrow-array/54.1.0/src/arrow_array/builder/generic_bytes_view_builder.rs.html#120-122, so this change might not help

Actually it helps. I have added a new test case test_sort_spill_utf8_strings. It will fail after removing the shrink_to_fit calls.

Here is where the 2X buffer growth come from:

1. sort_batch calls take_arrays, which calls take_bytes for string columns
2. take_bytes allocates a MutableBuffer for storing strings taken from the input array
3. take_bytes calls the extend_from_slice method of the values mutable buffer to append strings to the buffer, which in turn calls reserve to grow its space
4. reserve grows the size exponentially by a factor of 2

I can get the high-level idea of the more conservative memory accounting in point 3, it is used to also account for merged batches, but I get lost in the memory estimation details in the implementation (especially is this 2X estimation amplification only for merged batches, or also intermediate Rows), could you explain with a concrete example how everything is calculated? (for example, there is a ExternalSorter with 100M memory budget, and it will consume 10 batches, each with 20M size, how memory estimation is calculated in each step)

The 2X amplification is mainly for intermediate Rows, not for merged batches.

Let's assume that each batch is 10 MB, and we have 100 MB memory budget. The following diagram shows how the memory consumption become 100MB when performing merging.

Here is the detailed explanation:

1. We reserve 2X memory for each batch on insertion, so when in_mem_batches holds 5 batches and consumes 50 MB memory, we have already reserved 100 MB memory. The next insertion will trigger a merge, and possibly a spill.
2. When merge happens, each batch in in_mem_batches was sorted individually using sort_batch and fed into StreamingMergeBuilder to build a SortPreservingMergeStream. The batches were taken away from in_mem_batches, and the original batches will be dropped immediately after retrieving a sorted batch. We assume that the sorted batches has the same size as the original batches.
3. SortPreservingMergeStream polls one batch from each sorted stream, create a row representation or sorted array representation for each batch. The sorted batches were saved into in_progress and the row/sorted array representation were saved into cursors. We assume that the row representation or sorted array has the same size as the sorted batch. Now we have consumed 100MB.
4. SortPreservingMergeStream produces merged batches. We can assume that the overall memory consumption remains unchanged during this process, and certainly we need to reserve memory for merged batches. Each time we poll a merged batche from SortPreservingMergeStream, we try reserving memory for it. If the reservation fails, all future merged batches polled from the merged stream will be directly written to the spill file.

BTW, the 2X amplification for intermediate Rows seems to be conservative, but it is still not enough for some of the queries. The following query creates Rows larger than the original batches, and the query will fail when reserving memory for Rows:

-- sort-tpch Q5: 3 sort keys {(INTEGER, 7), (BIGINT, 10k), (BIGINT, 1.5M)} + no payload column

SELECT l_linenumber, l_suppkey, l_orderkey
FROM lineitem
ORDER BY l_linenumber, l_suppkey, l_orderkey

This can be worked around by configuring a larger sort_spill_reservation_bytes.

I had another interesting observation: spilling sort can be faster than memory unbounded sort in datafusion.

I tried running sort-tpch Q3 using this PR with #14642 cherry-picked onto it, and configured parquet.schema_force_view_types = false to mitigate #12136 (comment). Here are the test results obtained on a cloud instance with Intel(R) Xeon(R) Platinum 8269CY CPU @ 2.50GHz CPU:

$./target/release/dfbench sort-tpch --iterations 1 --path benchmarks/data/tpch_sf10 --memory-limit 1000M -q 3 -n1
Q3 iteration 0 took 93339.0 ms and returned 59986052 rows
Q3 avg time: 93339.00 ms
$./target/release/dfbench sort-tpch --iterations 1 --path benchmarks/data/tpch_sf10 --memory-limit 500M -q 3 -n1
Q3 iteration 0 took 81831.2 ms and returned 59986052 rows
Q3 avg time: 81831.18 ms
$./target/release/dfbench sort-tpch --iterations 1 --path benchmarks/data/tpch_sf10 --memory-limit 200M -q 3 -n1
Q3 iteration 0 took 77046.4 ms and returned 59986052 rows
Q3 avg time: 77046.36 ms
$./target/release/dfbench sort-tpch --iterations 1 --path benchmarks/data/tpch_sf10 -q 3 -n1
Q3 iteration 0 took 170416.1 ms and returned 59986052 rows
Q3 avg time: 170416.10 ms

When running without memory limit, we are merging tons of small sorted streams, this seems to be bad for performance. Memory limit enforces us to do merging before ingesting all the batches, so we are doing several smaller merges first and do a final merge at last to produce the result set. Coalescing batches into larger streams before merging seems to be a good idea.

Not directly related to the point of this PR but regarding
I had a hard time making DataFusion Comet work on cloud instances with 4GB memory per CPU core, partially because DataFusion is very likely to allocates more memory than reserved and run into OOM, or run into various kinds of memory reservation failures

I hope this apache/datafusion-comet#1369 helps. This may cause more spilling but I expect the chance of OOM is reduced. Still accurate tracking of memory is important, otherwise mis-tracking still causes OOM

use arrow::array::{RecordBatch, StringBuilder};
use arrow_schema::{DataType, Field, Schema};
use datafusion::execution::disk_manager::DiskManagerConfig;
use datafusion::execution::memory_pool::FairSpillPool;
use datafusion::execution::runtime_env::RuntimeEnvBuilder;
use datafusion::parquet::arrow::ArrowWriter;
use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext};
use futures::TryStreamExt;
use std::sync::Arc;

#[tokio::main(flavor = "multi_thread", worker_threads = 1)]
pub async fn main() {
    build_parquet();

    let env = RuntimeEnvBuilder::new()
        .with_disk_manager(DiskManagerConfig::default())
        .with_memory_pool(Arc::new(FairSpillPool::new(100 * 1024 * 1024)))
        .build_arc()
        .unwrap();

    let mut config = SessionConfig::new().with_sort_spill_reservation_bytes(32 * 1024 * 1024);
    config.options_mut().execution.parquet.schema_force_view_types = false;

    let ctx = SessionContext::new_with_config_rt(config, env);

    ctx.register_parquet(
        "big_strings",
        "/tmp/big_strings.parquet",
        ParquetReadOptions::default(),
    )
        .await
        .unwrap();

    let sql = "SELECT * FROM big_strings ORDER BY strings";
    println!("Sorting strings");
    ctx.sql(sql)
        .await
        .unwrap()
        .execute_stream()
        .await
        .unwrap()
        .try_for_each(|_| std::future::ready(Ok(())))
        .await
        .unwrap();
}

fn build_parquet() {
    if std::fs::File::open("/tmp/big_strings.parquet").is_ok() {
        println!("Using existing file at /tmp/big_strings.parquet");
        return;
    }
    println!("Generating test file at /tmp/big_strings.parquet");
    let file = std::fs::File::create("/tmp/big_strings.parquet").unwrap();
    let schema = Arc::new(Schema::new(vec![Field::new(
        "strings",
        DataType::Utf8,
        false,
    )]));
    let mut writer = ArrowWriter::try_new(file, schema.clone(), None).unwrap();

    for batch_idx in 0..100 {
        println!("Generating batch {} of 100", batch_idx);
        let mut string_array_builder =
            StringBuilder::with_capacity(1024 * 1024, 1024 * 1024 * 3 * 14);
        for i in 0..(1024 * 1024) {
            string_array_builder
                .append_value(format!("string-{}string-{}string-{}", i, i, i));
        }
        let array = Arc::new(string_array_builder.finish());
        let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
        writer.write(&batch).unwrap();
    }
    writer.close().unwrap();
}

called `Result::unwrap()` on an `Err` value: ResourcesExhausted("Failed to allocate additional 353536 bytes for ExternalSorterMerge[1] with 22948928 bytes already allocated for this reservation - 127190 bytes remain available for the total pool")

Thank you @kazuyukitanimura for the PR, i applied the PR try to fix the testing, but the above testing is still failed for me, i am not sure if i am missing something.

Thank you @kazuyukitanimura for the PR, i applied the PR try to fix the testing, but the above testing is still failed for me, i am not sure if i am missing something.

There are 2 problems:

1. The DataSourceExec may have many partitions, and each SortExec on that partition will only get a fair share of the 100MB pool, so each partition won't get enough memory to operate. This is still the case even with worker_threads = 1. If you still want to sort 4.2GB parquet file using 100MBs of memory, you can set .with_target_partitions(1) in your session config.
2. 100MB is not enough for the final merging with spill-reads. There will be roughly 200 spill files generated after ingesting all the batches, the size of a typical batch for this workload is 352380 bytes. The memory needed for merging will be 200 * (352380 bytes) * 2 > 100MB. Merging phase is unspillable so it requires a minimum amount of memory to operate. Raising the memory limit to 200MB will work for this particular workload.

One possible fix for problem 2 is to use a smaller batch size when writing batches to spill files, so that the unspillable memory required for the final spill-read merging will be smaller. Or we simply leave this problem as is and requires the user to raise the memory limit.

Thank you @kazuyukitanimura for the PR, i applied the PR try to fix the testing, but the above testing is still failed for me, i am not sure if i am missing something.

There are 2 problems:

1. The DataSourceExec may have many partitions, and each SortExec on that partition will only get a fair share of the 100MB pool, so each partition won't get enough memory to operate. This is still the case even with worker_threads = 1. If you still want to sort 4.2GB parquet file using 100MBs of memory, you can set .with_target_partitions(1) in your session config.
2. 100MB is not enough for the final merging with spill-reads. There will be roughly 200 spill files generated after ingesting all the batches, the size of a typical batch for this workload is 352380 bytes. The memory needed for merging will be 200 * (352380 bytes) * 2 > 100MB. Merging phase is unspillable so it requires a minimum amount of memory to operate. Raising the memory limit to 200MB will work for this particular workload.

One possible fix for problem 2 is to use a smaller batch size when writing batches to spill files, so that the unspillable memory required for the final spill-read merging will be smaller. Or we simply leave this problem as is and requires the user to raise the memory limit.

Thank you @Kontinuation for good explain, it makes sense to me, i will try it.

And for the problem 2, is it possible we introduce a spillable merging phase, will it be more safe?

Thank you @kazuyukitanimura for the PR, i applied the PR try to fix the testing, but the above testing is still failed for me, i am not sure if i am missing something.

There are 2 problems:

1. The DataSourceExec may have many partitions, and each SortExec on that partition will only get a fair share of the 100MB pool, so each partition won't get enough memory to operate. This is still the case even with worker_threads = 1. If you still want to sort 4.2GB parquet file using 100MBs of memory, you can set .with_target_partitions(1) in your session config.
2. 100MB is not enough for the final merging with spill-reads. There will be roughly 200 spill files generated after ingesting all the batches, the size of a typical batch for this workload is 352380 bytes. The memory needed for merging will be 200 * (352380 bytes) * 2 > 100MB. Merging phase is unspillable so it requires a minimum amount of memory to operate. Raising the memory limit to 200MB will work for this particular workload.

One possible fix for problem 2 is to use a smaller batch size when writing batches to spill files, so that the unspillable memory required for the final spill-read merging will be smaller. Or we simply leave this problem as is and requires the user to raise the memory limit.

Updated, it works after change to 1 partition and increase the memory limit.

Here is the detailed explanation:

1. We reserve 2X memory for each batch on insertion, so when in_mem_batches holds 5 batches and consumes 50 MB memory, we have already reserved 100 MB memory. The next insertion will trigger a merge, and possibly a spill.
2. When merge happens, each batch in in_mem_batches was sorted individually using sort_batch and fed into StreamingMergeBuilder to build a SortPreservingMergeStream. The batches were taken away from in_mem_batches, and the original batches will be dropped immediately after retrieving a sorted batch. We assume that the sorted batches has the same size as the original batches.
3. SortPreservingMergeStream polls one batch from each sorted stream, create a row representation or sorted array representation for each batch. The sorted batches were saved into in_progress and the row/sorted array representation were saved into cursors. We assume that the row representation or sorted array has the same size as the sorted batch. Now we have consumed 100MB.
4. SortPreservingMergeStream produces merged batches. We can assume that the overall memory consumption remains unchanged during this process, and certainly we need to reserve memory for merged batches. Each time we poll a merged batche from SortPreservingMergeStream, we try reserving memory for it. If the reservation fails, all future merged batches polled from the merged stream will be directly written to the spill file.

Thanks for the nice diagram, this explanation is super clear

Regarding point 3, I thought of a edge case can cause Rows way larger than 1X:
During merging, now only order key will be converted to row format.

• If select c1,c2 from t order by c1, c2, all columns will be convereted to Row, the estimation is close to 1X
• If select c1, ... c10 from t order by c1, only 1 column will be converted, the extra Row overhead is 0.1X. (I think this is okay to estimate more to be conservative)

Edge case: let's say input is a deduplicated StringViewArray (like a 10k rows batch with only 100 distinct values, but payload content are stored without duplication, the array elements are just referencing to the payload range), after converting to Row format, every row will be materialized, then the Row format will have 100X expansion
I think we need some mechanism to deal with this kind of edge case, perhaps this also applies to dictionary representation

For point 4, are the memory budget to hold merged batches come from sort_spill_reservation_bytes? Small sorted runs, and converted rows should have taken up all memory spaces at this stage.

Edge case: let's say input is a deduplicated StringViewArray (like a 10k rows batch with only 100 distinct values, but payload content are stored without duplication, the array elements are just referencing to the payload range), after converting to Row format, every row will be materialized, then the Row format will have 100X expansion I think we need some mechanism to deal with this kind of edge case, perhaps this also applies to dictionary representation

I agree that the current implementation uses a very rough estimation, and it could be way off from the actual memory consumption.

A better approach is to sort and generate the row representation of the batch right after we ingesting it, then we would know the exact size of sorted batches and their row representations held in memory. The merge phase for handling spilling could simply take away these data and perform merging without reserving more memory. However, this conflicts between some of the optimizations we did in the past:

• Parallel merge sort (#6162) #6308: sort are performed concurrently right before merging.

For point 4, are the memory budget to hold merged batches come from sort_spill_reservation_bytes? Small sorted runs, and converted rows should have taken up all memory spaces at this stage.

Yes. It may come from sort_spill_reservation_bytes, or the reduced memory usage after per-batch sorting because of the fetch option.

This PR can't solve
select * from t consumes 2GB memory to fully materialize the output, in the best case select * from t order by c1 consumes around 2GB memory, given sorting is an O(1) space operation. The current situation is it consumes 4GB+ (more than 2 times).

I think this is by design. The physical operators should not reserve memory for the batches it produces. It is the parent operator's duty to reserve memory for the batches fetched from children operators if it needs to hold the batches in memory for a long time.

Problem
I tried one query and this PR is not working as expected, I specified one query to run under 5GB memory (select * without order requires 7GB) but it's still consuming around 10GB, could you double check? I suspect we missed some details.

It is a problem of datafusion-cli. If datafusion-cli decides to hold all the result batches in memory, it should create a memory consumer for itself and reserve memory for the result batches.

datafusion-cli calls collect to hold the entire result set in memory before displaying it, this is unnecessary when maxrows is not unlimited. I've tried the following code to replace the collect call, and the maximum resident set size has reduced to 4.8 GB:

// Bounded stream; collected results are printed after all input consumed.
let schema = physical_plan.schema();
let mut stream = execute_stream(physical_plan, task_ctx.clone())?;
let mut results = vec![];
while let Some(batch) = stream.next().await {
    let batch = batch?;
    results.push(batch);
    if let MaxRows::Limited(max_rows) = print_options.maxrows {
        if results.len() >= max_rows {
            break;
        }
    }
}
adjusted.into_inner().print_batches(schema, &results, now)?;

Sorry, I mistakenly edited your original reply @Kontinuation, I'm trying to revert it back.

It is a problem of datafusion-cli. If datafusion-cli decides to hold all the result batches in memory, it should create a memory consumer for itself and reserve memory for the result batches.

Yes, you are right and the result is good. I made a mistake about the expected behavior: it should be measured RSS < memory limit + result size. I've tested several smaller memory limit and the result is correct.

This PR can't solve
select * from t consumes 2GB memory to fully materialize the output, in the best case select * from t order by c1 consumes around 2GB memory, given sorting is an O(1) space operation. The current situation is it consumes 4GB+ (more than 2 times).

I think this is by design. The physical operators should not reserve memory for the batches it produces. It is the parent operator's duty to reserve memory for the batches fetched from children operators if it needs to hold the batches in memory for a long time.

Problem
I tried one query and this PR is not working as expected, I specified one query to run under 5GB memory (select * without order requires 7GB) but it's still consuming around 10GB, could you double check? I suspect we missed some details.

It is a problem of datafusion-cli. If datafusion-cli decides to hold all the result batches in memory, it should create a memory consumer for itself and reserve memory for the result batches.

datafusion-cli calls collect to hold the entire result set in memory before displaying it, this is unnecessary when maxrows is not unlimited. I've tried the following code to replace the collect call, and the maximum resident set size has reduced to 4.8 GB:

// Bounded stream; collected results are printed after all input consumed.
let schema = physical_plan.schema();
let mut stream = execute_stream(physical_plan, task_ctx.clone())?;
let mut results = vec![];
while let Some(batch) = stream.next().await {
    let batch = batch?;
    results.push(batch);
    if let MaxRows::Limited(max_rows) = print_options.maxrows {
        if results.len() >= max_rows {
            break;
        }
    }
}
adjusted.into_inner().print_batches(schema, &results, now)?;

Good finding, may be we can file a issue to improve the datafusion-cli memory usage and reservation.

Thanks for the ping and thanks for working on this! This is an important feature for us (for training secondary indices on string columns) so I'm very thankful to see the effort 😄

I tried the reproducer from #10073 on this branch (babe5cd) and wasn't able to get it to pass and so I agree it doesn't seem to address all issues. However, what you're describing does seem to address the problems that I was seeing and so I think it probably is making good progress.

I also tried setting session_config.options_mut().execution.parquet.schema_force_view_types = false; and various values of sort_spill_reservation_bytes but didn't have much luck.

Hi @westonpace , i think the problem is we need to setting the partition count and also to increase the memory limit also for your case:

1. setting the partition count to 1:
 .with_target_partitions(1)

2. increasing the memory limit for example to 300MB

I can share the code to pass for your case, also clean the build parquet file and retry:

use arrow::array::{RecordBatch, StringBuilder};
use arrow_schema::{DataType, Field, Schema};
use datafusion::execution::disk_manager::DiskManagerConfig;
use datafusion::execution::memory_pool::FairSpillPool;
use datafusion::execution::runtime_env::RuntimeEnvBuilder;
use datafusion::parquet::arrow::ArrowWriter;
use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext};
use futures::TryStreamExt;
use std::sync::Arc;

#[tokio::main(flavor = "multi_thread", worker_threads = 1)]
pub async fn main() {
    build_parquet();

    let env = RuntimeEnvBuilder::new()
        .with_disk_manager(DiskManagerConfig::default())
        .with_memory_pool(Arc::new(FairSpillPool::new(300 * 1024 * 1024)))
        .build_arc()
        .unwrap();

    let mut config = SessionConfig::new().with_sort_spill_reservation_bytes(32 * 1024 * 1024).with_target_partitions(1);
    config.options_mut().execution.parquet.schema_force_view_types = false;

    let ctx = SessionContext::new_with_config_rt(config, env);

    ctx.register_parquet(
        "big_strings",
        "/tmp/big_strings.parquet",
        ParquetReadOptions::default(),
    )
        .await
        .unwrap();

    let sql = "SELECT * FROM big_strings ORDER BY strings";
    println!("Sorting strings");
    ctx.sql(sql)
        .await
        .unwrap()
        .execute_stream()
        .await
        .unwrap()
        .try_for_each(|_| std::future::ready(Ok(())))
        .await
        .unwrap();
}

fn build_parquet() {
    if std::fs::File::open("/tmp/big_strings.parquet").is_ok() {
        println!("Using existing file at /tmp/big_strings.parquet");
        return;
    }
    println!("Generating test file at /tmp/big_strings.parquet");
    let file = std::fs::File::create("/tmp/big_strings.parquet").unwrap();
    let schema = Arc::new(Schema::new(vec![Field::new(
        "strings",
        DataType::Utf8,
        false,
    )]));
    let mut writer = ArrowWriter::try_new(file, schema.clone(), None).unwrap();

    for batch_idx in 0..100 {
        println!("Generating batch {} of 100", batch_idx);
        let mut string_array_builder =
            StringBuilder::with_capacity(1024 * 1024, 1024 * 1024 * 3 * 14);
        for i in 0..(1024 * 1024) {
            string_array_builder
                .append_value(format!("string-{}string-{}string-{}", i, i, i));
        }
        let array = Arc::new(string_array_builder.finish());
        let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
        writer.write(&batch).unwrap();
    }
    writer.close().unwrap();
}

Details reason here:

#14644 (comment)

1. The DataSourceExec may have many partitions, and each SortExec on that partition will only get a fair share of the 100MB pool, so each partition won't get enough memory to operate. This is still the case even with worker_threads = 1. If you still want to sort 4.2GB parquet file using 100MBs of memory, you can set .with_target_partitions(1) in your session config.

2. 100MB is not enough for the final merging with spill-reads. There will be roughly 200 spill files generated after ingesting all the batches, the size of a typical batch for this workload is 352380 bytes. The memory needed for merging will be 200 * (352380 bytes) * 2 > 100MB. Merging phase is unspillable so it requires a minimum amount of memory to operate. Raising the memory limit to 200MB will work for this particular workload.

And for future work we can have a complete spill solution tracking here:

#14692

This PR can't solve
select * from t consumes 2GB memory to fully materialize the output, in the best case select * from t order by c1 consumes around 2GB memory, given sorting is an O(1) space operation. The current situation is it consumes 4GB+ (more than 2 times).

I think this is by design. The physical operators should not reserve memory for the batches it produces. It is the parent operator's duty to reserve memory for the batches fetched from children operators if it needs to hold the batches in memory for a long time.

This is a great observation and I think holds true in practice

However I don't think it is reflected in the documention here
https://docs.rs/datafusion/latest/datafusion/execution/memory_pool/trait.MemoryPool.html

Would someone be willing to update those docs?

Problem
I tried one query and this PR is not working as expected, I specified one query to run under 5GB memory (select * without order requires 7GB) but it's still consuming around 10GB, could you double check? I suspect we missed some details.

It is a problem of datafusion-cli. If datafusion-cli decides to hold all the result batches in memory, it should create a memory consumer for itself and reserve memory for the result batches.

datafusion-cli calls collect to hold the entire result set in memory before displaying it, this is unnecessary when maxrows is not unlimited. I've tried the following code to replace the collect call, and the maximum resident set size has reduced to 4.8 GB:

This is also another great observation that it would be great to get captured in a ticket

Thanks everyone -- this is a great example of collaboration and what the collective efforts of many contributors can do. I love it.

100MB is not enough for the final merging with spill-reads. There will be roughly 200 spill files generated after ingesting all the batches, the size of a typical batch for this workload is 352380 bytes. The memory needed for merging will be 200 * (352380 bytes) * 2 > 100MB. Merging phase is unspillable so it requires a minimum amount of memory to operate. Raising the memory limit to 200MB will work for this particular workload.

Thank you for this. You are correct, with this addition then my test passes. #14692 sounds awesome but I think we can also go far with these current fixes in the meantime 🚀

Problem
I tried one query and this PR is not working as expected, I specified one query to run under 5GB memory (select * without order requires 7GB) but it's still consuming around 10GB, could you double check? I suspect we missed some details.

It is a problem of datafusion-cli. If datafusion-cli decides to hold all the result batches in memory, it should create a memory consumer for itself and reserve memory for the result batches.
datafusion-cli calls collect to hold the entire result set in memory before displaying it, this is unnecessary when maxrows is not unlimited. I've tried the following code to replace the collect call, and the maximum resident set size has reduced to 4.8 GB:

This is also another great observation that it would be great to get captured in a ticket

Created a follow-up for this improvement:
#14751